#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
为大乐透创建基础和增强版LSTM-CRF模型文件的脚本
"""

import os
import sys
import torch
import torch.nn as nn
from datetime import datetime

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, '..', '..'))
sys.path.append(project_root)

class SimpleLstmCRFModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, output_seq_length, num_layers=1):
        super(SimpleLstmCRFModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.output_seq_length = output_seq_length
        
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_seq_length * output_dim)
        
    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        fc_out = self.fc(lstm_out[:, -1, :])
        fc_out = fc_out.view(-1, self.output_seq_length, self.output_dim)
        return fc_out

class EnhancedLstmCRFModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, output_seq_length, num_layers=2):
        super(EnhancedLstmCRFModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.output_seq_length = output_seq_length
        
        # 使用两层LSTM以增强模型能力
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=num_layers, batch_first=True, dropout=0.1)
        self.fc = nn.Linear(hidden_dim, output_seq_length * output_dim)
        self.dropout = nn.Dropout(0.1)
        
    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        lstm_out = self.dropout(lstm_out)
        fc_out = self.fc(lstm_out[:, -1, :])
        fc_out = fc_out.view(-1, self.output_seq_length, self.output_dim)
        return fc_out

def create_dlt_models():
    """创建大乐透基础和增强版LSTM-CRF模型"""
    print("[START] 开始创建大乐透LSTM-CRF模型")
    print("=" * 50)
    
    # 配置路径
    base_model_path = os.path.join(project_root, 'scripts', 'dlt', 'dlt_model.pth')
    enhanced_model_path = os.path.join(project_root, 'scripts', 'dlt', 'enhanced_lstm_crf_model.pth')
    
    print(f"[SAVE] 基础模型保存路径: {base_model_path}")
    print(f"[SAVE] 增强版模型保存路径: {enhanced_model_path}")
    
    try:
        # 创建目录（如果不存在）
        os.makedirs(os.path.dirname(base_model_path), exist_ok=True)
        
        # 初始化模型参数（大乐透）
        input_dim = 7
        hidden_dim = 128
        red_output_dim = 35  # 红球数量
        blue_output_dim = 12  # 蓝球数量
        red_output_seq_length = 5  # 红球序列长度
        blue_output_seq_length = 2  # 蓝球序列长度
        
        # 创建基础LSTM-CRF模型
        print("创建基础LSTM-CRF模型...")
        red_model = SimpleLstmCRFModel(input_dim, hidden_dim, red_output_dim, red_output_seq_length)
        blue_model = SimpleLstmCRFModel(input_dim, hidden_dim, blue_output_dim, blue_output_seq_length)
        
        # 保存基础模型
        torch.save({
            'red_model': red_model.state_dict(),
            'blue_model': blue_model.state_dict(),
            'input_dim': input_dim,
            'hidden_dim': hidden_dim,
            'red_output_dim': red_output_dim,
            'blue_output_dim': blue_output_dim,
            'red_output_seq_length': red_output_seq_length,
            'blue_output_seq_length': blue_output_seq_length,
            'model_type': 'dlt_lstm_crf',
            'created_time': datetime.now().isoformat()
        }, base_model_path)
        
        print("[OK] 基础LSTM-CRF模型创建成功！")
        
        # 创建增强版LSTM-CRF模型
        print("创建增强版LSTM-CRF模型...")
        red_enhanced_model = EnhancedLstmCRFModel(input_dim, hidden_dim, red_output_dim, red_output_seq_length, num_layers=2)
        blue_enhanced_model = EnhancedLstmCRFModel(input_dim, hidden_dim, blue_output_dim, blue_output_seq_length, num_layers=2)
        
        # 保存增强版模型
        torch.save({
            'red_lstm_model': red_enhanced_model.state_dict(),
            'blue_lstm_model': blue_enhanced_model.state_dict(),
            'input_dim': input_dim,
            'hidden_dim': hidden_dim,
            'red_output_dim': red_output_dim,
            'blue_output_dim': blue_output_dim,
            'red_output_seq_length': red_output_seq_length,
            'blue_output_seq_length': blue_output_seq_length,
            'model_type': 'dlt_enhanced_lstm_crf',
            'lottery_type': 'dlt',
            'created_time': datetime.now().isoformat()
        }, enhanced_model_path)
        
        print("[OK] 增强版LSTM-CRF模型创建成功！")
        
        print(f"[INFO] 基础模型已保存到: {base_model_path}")
        print(f"[INFO] 增强版模型已保存到: {enhanced_model_path}")
        return True
        
    except Exception as e:
        print(f"[ERROR] 创建模型过程出错: {e}")
        import traceback
        print(traceback.format_exc())
        return False

def main():
    """主函数"""
    print("[TARGET] 大乐透LSTM-CRF模型创建程序")
    print(f"[TIME] 开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print()
    
    success = create_dlt_models()
    
    print()
    if success:
        print("[SUCCESS] 所有模型创建完成！")
    else:
        print("[FAIL] 模型创建失败！")
    
    print(f"[TIME] 结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

if __name__ == "__main__":
    main()